package com.hp.test.interview.misc;

import com.hp.test.interview.graph.BinaryMinHeap;

/**
 * Kth largest in rowwise and column wise sorted array
 * http://www.geeksforgeeks.org/kth-smallest-element-in-a-row-wise-and-column-wise-sorted-2d-array-set-1/
 */
public class KthLargestInRowiseColumnWiseSorted2DArray {

    public static void main(String args[]) {
        int input[][] = {
                {1, 5, 11, 21}, {6, 8, 12, 22}, {7, 9, 15, 26}, {11, 13, 18, 28}
        };
        KthLargestInRowiseColumnWiseSorted2DArray kl = new KthLargestInRowiseColumnWiseSorted2DArray();
        System.out.println(kl.kthLargest(input, 10));
    }

    public int kthLargest(int input[][], int k) {

        BinaryMinHeap<Integer> minHeap = new BinaryMinHeap<Integer>();
        int c = input[0].length;
        int total = input.length * c;
        minHeap.add(input[0][0], 0);
        for (int i = 1; i < k; i++) {
            int minIndex = minHeap.extractMin();
            int minRow = minIndex / c;
            int minCol = minIndex % c;

            int downNeighbor = (minRow + 1) * c + minCol;
            int rightNeighbor;
            if (minCol == (c - 1)) {
                rightNeighbor = total;
            } else {
                rightNeighbor = minRow * c + (minCol + 1);
            }
            if (downNeighbor < total && !minHeap.containsData(downNeighbor)) {
                minHeap.add(input[minRow + 1][minCol], downNeighbor);
            }

            if (rightNeighbor < total && !minHeap.containsData(rightNeighbor)) {
                minHeap.add(input[minRow][minCol + 1], rightNeighbor);
            }
        }

        int minIndex = minHeap.extractMin();
        return input[minIndex / c][minIndex % c];
    }
}
